import os
import cv2
import numpy as np

def grabcut_segmentation(img, rect):
    mask = np.zeros(img.shape[:2],np.uint8)
    bg_model = np.zeros((1,65),np.float64)
    fg_model = np.zeros((1,65),np.float64)
    cv2.grabCut(img, mask, rect, bg_model, fg_model, 5, cv2.GC_INIT_WITH_RECT)
    mask2 = np.where((mask==2)|(mask==0),0,1).astype('uint8')
    img = img*mask2[:,:,np.newaxis]
    return img

# 文件夹路径
folder_path = r'D:\kaggle\final\single'
output_folder_path = r'D:\kaggle\final\segmentation'

# 检查输出文件夹是否存在，如果不存在则创建
if not os.path.exists(output_folder_path):
    os.makedirs(output_folder_path)

# 获取文件夹中的所有文件
files = os.listdir(folder_path)

# 对每一个文件进行处理
for i, file in enumerate(files):

    img_path = os.path.join(folder_path, file)
    img = cv2.imread(img_path)

    rect = (60,83,400,380)
    segmented_img = grabcut_segmentation(img, rect)

    output_img_path = os.path.join(output_folder_path, f'{i + 1}.jpg')
    cv2.imwrite(output_img_path, segmented_img)